#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Software: PyCharm
# @Version : Python-
# @Author  : Shengji He
# @Email   : hsjbit@163.com
# @File    : AverageLevelsBinaryTree.py
# @Time    : 2020/9/12 10:25
# @Description:
from typing import List
from collections import deque


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


class Solution:
    def averageOfLevels(self, root: TreeNode) -> List[float]:
        """
        Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.
        Example 1:
            Input:
                3
               / \
              9  20
                /  \
               15   7
            Output: [3, 14.5, 11]
            Explanation:
                The average value of nodes on level 0 is 3,  on level 1 is 14.5, and on level 2 is 11.
                Hence return [3, 14.5, 11].
        Note:
            - The range of node's value is in the range of 32-bit signed integer.

        :param root:
        :return:
        """
        counts = []
        totals = []
        def dfs(root: TreeNode, level: int):
            if not root:
                return
            if level < len(totals):
                totals[level] += root.val
                counts[level] += 1
            else:
                totals.append(root.val)
                counts.append(1)
            dfs(root.left, level + 1)
            dfs(root.right, level + 1)
        dfs(root, 0)
        return [value / count for value, count in zip(totals, counts)]

    def averageOfLevels_ways2(self, root: TreeNode) -> List[float]:
        ans = []
        queue = deque([root])
        while queue:
            total = 0
            counts = len(queue)
            for _ in range(counts):
                node = queue.popleft()
                total += node.val
                if node.left:
                    queue.append(node.left)
                if node.right:
                    queue.append(node.right)
            ans.append(total / counts)
        return ans


if __name__ == '__main__':
    S = Solution()
    root = TreeNode(3)
    root.left = TreeNode(9)
    root.right = TreeNode(20)
    root.right.left = TreeNode(15)
    root.right.right = TreeNode(7)
    print(S.averageOfLevels(root))
    print(S.averageOfLevels_ways2(root))
    print('done')
